'''
Created on Sep 2, 2009

@author: mkiyer
'''

from veggie.io.interval import read_bed_file
from veggie.sample.samplegroup import SampleGroup
from optparse import OptionParser
import h5py
import logging
import matplotlib.pyplot as plt
import numpy as np
import veggie.visualization.plots as oncoplot
import sys

def get_sample_coverages(chrom, start, end, sample_groups, coverage_hdf5):
    sample_coverages = {}
    for sgroup in sample_groups:
        for s in sgroup.samples:
            covarray = coverage_hdf5[s][chrom][start:end]
            sample_coverages[s] = covarray
    return sample_coverages

def plot_expr(chrom, start, end, sample_groups, coverage_hdf5,
              **kwargs):
    '''
    kwargs:
    TODO
    '''    
    # compute coverage
    sample_coverages = get_sample_coverages(chrom, start, end, sample_groups, coverage_hdf5)
    summed_coverages = {}
    for s, covarray in sample_coverages.iteritems():
        summed_coverages[s] = 1000.0 * (np.sum(covarray) / float(end - start))
    
    # setup figure
    default_suptitle = 'Expression for %s:%d-%d' % (chrom, start, end)
    default_figsize = (12,8)
    #default_window_title = 'Boxplot'
    
    fig = plt.figure(figsize=kwargs.get('figsize', default_figsize))
    fig.suptitle(kwargs.get('title', default_suptitle))
    #fig.canvas.set_window_title(kwargs.get('window_title', default_window_title))
    #fig.subplots_adjust(left=0.075, right=0.95, top=0.9, bottom=0.25)

    ax1 = fig.add_subplot(221)
    ax2 = fig.add_subplot(223)
    ax3 = fig.add_subplot(122)
    # Top plot shows individual samples
    oncoplot.plot_sample_coverage(chrom, start, end, 
                                  sample_groups, sample_coverages,
                                  axes=ax1)
    # Bottom plot shows average for groups
    oncoplot.plot_sample_group_average_coverage(chrom, start, end, 
                                                sample_groups, 
                                                sample_coverages,
                                                axes=ax2)
    # Boxplot
    oncoplot.plot_boxplot(chrom, start, end, 
                          sample_groups, summed_coverages,
                          axes=ax3)    
    plt.savefig('./plots/%s_%d-%d_expr.png' % (chrom, start, end))
    plt.close()
    return


if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG)

    optionparser = OptionParser("usage: %prog [options]")
    optionparser.add_option("-i", "--intervals", dest="bed_file")
    optionparser.add_option("-g", "--groups", dest="sample_groups")
    optionparser.add_option("-c", "--coverage", dest="coverage_file")    
    optionparser.add_option("-o", "--output", dest="outfile",
                            default=None,
                            help="output file [default: %default]")
    (options, args) = optionparser.parse_args()
    if options.outfile == None:
        outfhd = sys.stdout
    else:
        outfhd = open(options.outfile, 'w')
    
    # intervals - bed file
    intervals = read_bed_file(open(options.bed_file))
    # sample groups
    # TODO: load from XML file
    t = SampleGroup(name='ETS- Tumor', color='blue', cond={'sample_type': 'Tissue', 
                                                           'tissue_type': 'Prostate',
                                                           'diagnosis': 'Localized Tissue', 
                                                           'treated':'FALSE', 
                                                           'ets':'ETS-'})
    c = SampleGroup(name='Benign', color='green', cond={'sample_type': 'Tissue', 
                                                        'tissue_type': 'Prostate', 
                                                        'diagnosis': 'Benign Tissue',
                                                        'treated':'FALSE'})
    sample_groups = [c,t]
    # open the coverage database file
    h5f = h5py.File(options.coverage_file, 'r')
    
    # only consider samples with coverage data
    for sgroup in sample_groups:
        filtered_sgroup = []
        for s in sgroup.samples:
            if s in h5f:
                logging.debug('    added treatment: %s' % s)
                filtered_sgroup.append(s)
            else:
                logging.warning('    %s coverage does not exist yet.. skipping' % s)
        sgroup.samples = filtered_sgroup             

    # process intervals
    for interval in intervals:
        chrom, start, end = interval
        logging.debug('plotting %s:%d-%d' % (chrom, start, end))
        plot_expr(chrom, start, end, sample_groups, h5f)
    # close
    h5f.close()